#include <iostream>
#include <cmath>
#include <cuda_runtime.h>
#include <vector>
#include <chrono>

void cpuGEMM(const std::vector<float>& A,
             const std::vector<float>& B,
             std::vector<float>& C,
             int M, int N, int K) {
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float value = 0.0f;
            for (int k = 0; k < K; ++k) {
                value += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = value;
        }
    }
}

__global__ void wmmaGEMM(float* A, float* B, float* C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M && col < N) {
        float value = 0;
        for (int k = 0; k < K; ++k) {
            value += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = value;
    }
}

bool compareResults(const std::vector<float>& ref,
                    const std::vector<float>& gpu,
                    int M, int N) {
    bool flag = true;
    const float epsilon = 1.0f;
    for (int i = 0; i < M * N; ++i) {
        if (std::fabs(ref[i] - gpu[i]) > epsilon) {
            std::cout << "Mismatch at index " << i << ": CPU=" << ref[i] << ", GPU=" << gpu[i] << "\n";
            std::cout << "CPU value: " << ref[i] << ", GPU value: " << gpu[i] << "\n";
            std::cout << "Difference: " << std::fabs(ref[i] - gpu[i]) << "\n";
            flag = false;
        }
    }
    return flag;
}

int main() {
    int M = 1024;
    int N = 1024;
    int K = 1024;
    int BLK_M = 32;
    int BLK_N = 32;

    size_t sizeA = M * K * sizeof(float);
    size_t sizeB = K * N * sizeof(float);
    size_t sizeC = M * N * sizeof(float);

    std::vector<float> h_A(M * K, 1.0f);
    std::vector<float> h_B(K * N, 1.0f);
    std::vector<float> h_C_cpu(M * N, 0.0f);
    std::vector<float> h_C_gpu(M * N, 0.0f);

    for (int i = 0; i < M * K; ++i) h_A[i] = static_cast<float>(i % 1024);
    for (int i = 0; i < K * N; ++i) h_B[i] = static_cast<float>(i % 1024);

    std::cout << "Computing reference result on CPU...\n";
    auto cpu_start = std::chrono::high_resolution_clock::now();
    cpuGEMM(h_A, h_B, h_C_cpu, M, N, K);
    auto cpu_end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> cpu_time = cpu_end - cpu_start;
    std::cout << "Reference result computed.\n";

    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, sizeA);
    cudaMalloc(&d_B, sizeB);
    cudaMalloc(&d_C, sizeC);

    cudaMemcpy(d_A, h_A.data(), sizeA, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B.data(), sizeB, cudaMemcpyHostToDevice);
    cudaMemcpy(d_C, h_C_gpu.data(), sizeC, cudaMemcpyHostToDevice);

    dim3 block(BLK_N, BLK_M);
    dim3 grid((N+BLK_N-1)/BLK_N, (M+BLK_M-1)/BLK_M);

    std::cout << "Launching kernel...\n";

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    cudaEventRecord(start);

    wmmaGEMM<<<grid, block>>>(d_A, d_B, d_C, M, N, K);

    cudaEventRecord(stop);
    cudaEventSynchronize(stop);

    float gpu_time = 0.0f;
    cudaEventElapsedTime(&gpu_time, start, stop);

    cudaDeviceSynchronize();

    cudaMemcpy(h_C_gpu.data(), d_C, sizeC, cudaMemcpyDeviceToHost);

    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    if (compareResults(h_C_cpu, h_C_gpu, M, N)) {
        std::cout << "Results match!\n";
    } else {
        std::cout << "Results do not match!\n";
    }
    std::cout << "===========================TIME==========================\n";
    std::cout << "CPU Time: " << cpu_time.count() * 1000.0 << " ms\n";
    std::cout << "GPU Time: " << gpu_time << " ms\n";

    std::cout << "Done!\n";
}